import json
import time
import random
from Cryptodome.PublicKey import RSA
from base64 import b64encode
from Cryptodome.Signature import pkcs1_15
from Cryptodome.Hash import SHA256
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
import base64
import requests
from xml.etree import ElementTree as ET


# 需要request
class WechatPay(object):
    def __init__(self, request):
        """=========需要申请的参数========="""
        # 小程序的APPID
        self.APPID = ""
        # 微信支付商户号
        self.MChID = ""
        # 商户证书序列号
        self.serialNo = ""
        # API v3密钥
        self.APIv3_KEY = ""
        # 商户证书私钥
        self.PRIVATE_KEY = open("../1900009191_wxp_cert.pem").read()
        """=========需要申请的参数========="""
        # 请求参数
        self.request = request

        # 初始化获取prepay_id所需数据
        self.init_body()

    def init_body(self):
        # 请求体参数
        self.body = json.dumps({
            "appid": self.APPID,
            "mchid": self.MChID,
            "description": "xxx",
            "out_trade_no": self.request.data['out_trade_no'],
            "notify_url": "https://XXXXX/payNotify",  # 后端接收回调通知的接口
            "amount": {"total": self.request.data['price'] * 100, "currency": "CNY"},  # 正式上线price要*100，微信金额单位为分。
            "payer": {"openid": self.request.data['openid']},
        }).encode('utf8')
        # prepay_id
        self.prepay_id = ""
        # 解密后的证书
        self.certificate = ""
        # 生成时间戳
        self.time_stamps = self.timeStamps()
        # 随机12位字符串
        self.random_str = self.nonceStr()
        # 数据类型
        self.type = ""

    # 随机字符串，长度为32个字符以下
    def nonceStr(self):
        # 大写字母中生成随机12位字符串数据chr用于将65-90(阿斯克码)转为大写字母
        return "".join([chr(random.randint(65, 90)) for _ in range(12)])

    # 生成请求时间戳
    def timeStamps(self):
        return str(int(time.time()))

    # 构造的请求签名串
    @property
    def clean_body(self):
        if self.type == 'pid':
            return f"POST\n{'/v3/pay/transactions/jsapi'}\n{self.time_stamps}\n{self.random_str}\n{self.body}\n".encode(
                'utf8')
        if self.type == "verify_id":
            return f"GET\n{'/v3/certificates'}\n{self.time_stamps}\n{self.random_str}\n""\n"

        return f"{self.APPID}\n{self.time_stamps}\n{self.random_str}\nprepay_id={self.prepay_id}".encode('utf8')

    # 签名
    @property
    def get_sign(self):
        rsa_key = RSA.importKey(self.PRIVATE_KEY)
        signer = pkcs1_15.new(rsa_key)
        digest = SHA256.new(self.clean_body)
        sign = b64encode(signer.sign(digest)).decode('utf-8')
        self.type = ""
        return sign

    @property
    def create_headers(self):
        return {
            'Content-Type': 'application/json',
            'Accept': 'application/json',
            'User-Agent': '*/*',
            'Authorization': 'WECHATPAY2-SHA256-RSA2048 ' + f'mchid="{self.MChID}",nonce_str="{self.random_str}",signature="{self.get_sign}",timestamp="{self.time_stamps}",serial_no="{self.serialNo}"'
        }

    # 获取prepay_id
    def getPrePayId(self):
        url = 'https://api.mch.weixin.qq.com/v3/pay/transactions/jsapi'
        self.type = 'pid'
        # 10.发送请求获得prepay_id
        response = requests.post(url, data=self.body, headers=self.create_headers)  # 获取预支付交易会话标
        self.prepay_id = response.json()['prepay_id']
        return response

    # 验签
    def verify_main(self, response):
        url = "https://api.mch.weixin.qq.com/v3/certificates"
        # 11.应答签名验证
        wechatpaySerial = response.headers['Wechatpay-Serial']  # 获取HTTP头部中包括回调报文的证书序列号
        wechatpaySignature = response.headers['Wechatpay-Signature']  # 获取HTTP头部中包括回调报文的签名
        wechatpayTimestamp = response.headers['Wechatpay-Timestamp']  # 获取HTTP头部中包括回调报文的时间戳
        wechatpayNonce = response.headers['Wechatpay-Nonce']  # 获取HTTP头部中包括回调报文的随机串
        self.type = "verify_id"
        result = requests.get(url, headers=self.create_headers)
        cert = result.json()
        nonce = cert["data"][0]['encrypt_certificate']['nonce']
        ciphertext = cert["data"][0]['encrypt_certificate']['ciphertext']
        associated_data = cert["data"][0]['encrypt_certificate']['associated_data']
        serial_no = cert["data"][0]['serial_no']
        self.certificate = self.decrypt(nonce, ciphertext, associated_data)
        # 判断序列号是否相同
        if wechatpaySerial == serial_no:
            data = f"{wechatpayTimestamp}\n{wechatpayNonce}\n{response.text}\n"
            try:
                self.verify(data, wechatpaySignature)
                return True
            except (ValueError, TypeError):
                return False
        return False

    def verify(self, data, signature):
        key = RSA.importKey(self.certificate)  # 这里直接用了解密后的证书，但没有去导出公钥，似乎也是可以的。怎么导公钥还没搞懂。
        verifier = pkcs1_15.new(key)
        hash_obj = SHA256.new(data.encode('utf8'))
        return verifier.verify(hash_obj, base64.b64decode(signature))

    # AEAD_AES_256_GCM解密
    def decrypt(self, nonce, ciphertext, associated_data):
        # 转为字节类型
        key_bytes = str.encode(self.APIv3_KEY)
        nonce_bytes = str.encode(nonce)
        ad_bytes = str.encode(associated_data)
        # base64解码
        data = base64.b64decode(ciphertext)
        # AESGCM解密
        aesgcm = AESGCM(key_bytes)
        # 返回结果
        return aesgcm.decrypt(nonce_bytes, data, ad_bytes)

    def run(self):
        response = self.getPrePayId()
        # 验签
        if self.verify_main(response):
            data = {
                'timeStamp': self.time_stamps,
                'nonceStr': self.random_str,
                'package': f"prepay_id={self.prepay_id}",
                'signType': "RSA",
                'paySign': self.get_sign
            }
            return data
        return {'code': "204", "msg": "验签失败"}


if __name__ == '__main__':
    request = {'data': {}}
    pay = WechatPay(request)
    data_list = pay.run()
    # print(pay.timeStamps())
